from typing import Optional

from com.example.leetcode.tree.common import TreeUtils
from com.example.leetcode.tree.common.TreeUtils import TreeNode


class Solution:
    def deleteNode(self, root: Optional[TreeNode], key: int) -> Optional[TreeNode]:
        if root is None:
            return None
        if key < root.val:
            ## key更小，去左边
            root.left = self.deleteNode(root.left, key)
        elif key > root.val:
            ## key更大，去右边
            root.right = self.deleteNode(root.right, key)
        else:
            if root.left is None:
                return root.right
            elif root.right is None:
                return root.left
            else:
                min_node = self.get_min(root.right)
                root.val = min_node.val
                ## 更新右子树
                root.right = self.deleteNode(root.right, root.val)
        return root

    def get_min(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
        while root.left is not None:
            root = root.left
        return root


if __name__ == '__main__':
    s = Solution()
    nums = [5, 3, 6, 2, 4, None, 7]
    key = 3
    root = TreeUtils.generate_binary_tree(nums)
    TreeUtils.print_tree(root)

    result = s.deleteNode(root, key)
    TreeUtils.print_tree(result)
